#!/usr/bin/python
# -*- coding: utf-8 -*-

'''
IOTMQ API: Subscriber
'''
import urllib2, urllib
import logging
import json
from ConfigParser import ConfigParser
import zc.zk

class Subscribe():
    def __init__(self, conf_path):
        '''
        init subscibe api class
        '''
        try:
            self.conf = ConfigParser()
            self.conf.read(conf_path)
        except Exception, e:
            print 'read conf file failed, %s' % str(e)
            exit

        # init logger
        log_level = {}
        log_level['DEBUG'] = logging.DEBUG
        log_level['INFO'] = logging.INFO
        log_level['ERROR'] = logging.ERROR
        try:
            loglevel = log_level[self.conf.get('all', 'log-level')]
        except Exception, e:
            loglevel = logging.DEBUG

        formatstr = '%(asctime)s %(name)-8s %(levelname)-8s %(message)s'
        logging.basicConfig(level = loglevel, format = formatstr)
        self.logger = logging.getLogger('broker')
        self.logger.info('logger init OK.')

    def connect(self, user, passwd):
        '''
        connect to iot master server
        args:
            user: registered username
            passwd: passwd for username
        return:
            success, return True
            failed, return False
        '''
        meta_server = self.conf.get('all', 'meta-server')
        logging.info('find meta server address %s', meta_server)
        # init zk
        try:
            self.zk = zc.zk.ZooKeeper(meta_server)
        except Exception, e:
            logging.exception('connect to meta server failed, %s' % str(e))
            return
        # find master
        try:
            master = sorted(self.zk.children('/iotmq/master'))
            if not len(master):
                logging.error('no master running...')
                exit
            self.master = master[0]
            logging.info('find master %s', master[0])
        except Exception, e:
            logging.exception('find master error, %s' % str(e))
        # register node
        return True

    def subscribe(self, queue, subid, maxlen):
        '''
        subscribe message
        args:
            subid: subpoint, where sub start
            maxlen: max count message get per sub
        return:
            list, with sub message
        '''
        sub_id = subid
        while True:
            # find ds server
            master_uri = 'http://%s/%s/%d' % (self.master, queue, sub_id)
            logging.debug('master uri: %s', master_uri)

            block_info = None
            try:
                ret = urllib2.urlopen(master_uri).read()
                block_info = json.loads(ret)
            except Exception, e:
                logging.exception('connect to master failed, %s' % str(e))
                exit

            data_server = "http://%s/%s/%d/%d" % (block_info['ds_group'], queue, block_info['id'], sub_id)
            logging.debug('get data server uri: %s' % data_server)
            try :
                ret = urllib2.urlopen(data_server).read()
                logging.debug('read data: %s' % ret)
                return ret
            except Exception, e:
                logging.exception('get data failed, %s' % str(e))
                exit

    def disconnect(self):
        '''
        disconnect with server
        '''
        return True

if __name__ == "__main__":
    sub = Subscribe('./conf/pub.conf')
    sub.connect('root', 'root')
    counter = 0
    sub_id = 1
    while counter < 5000000:
        ret = sub.subscribe('test-queue', sub_id, 100)

        data = []
        if ret:
            data = json.loads(ret)
        sub_id = data[-1]['_id']
        logging.info('===> sub(%d) OK, count %d' % (sub_id, len(data)))
        counter += len(data)
